library(dplyr)
library(ggplot2)
library(viridis)
library(ggpubr)

if (!file.exists("output")) {
  dir.create("output")
}

# load data
data <- read.csv("Figure_7a_data.csv")

# Define function to generate plots from relative CN signature attributions
cn_sig_landscape <- function(sigs_rel) {
  signatures <- names(sigs_rel)[-1]
  # calculate frequency and mean relative attribution in tumors from each signature by subsite
  sigs_cat <- sigs_rel %>%
    tibble::column_to_rownames("donor_id") %>%
    mutate_if(is.numeric, ~ ifelse(. > 0, 1, 0))
  perc_sig <- sigs_rel %>%
    summarise_if(is.numeric, mean) %>%
    t()

  summary <- sigs_cat %>%
    summarise_if(is.numeric, sum) %>%
    t() %>%
    as.data.frame() %>%
    rename(n_pos = "V1") %>%
    mutate(freq = n_pos / nrow(sigs_rel)) %>%
    cbind(perc_sig) %>%
    tibble::rownames_to_column("sig")

  order <- summary[order(nrow(summary):1), ]

  ggplot(summary, aes(x = "HNC", y = sig, color = perc_sig, size = ifelse(freq == 0, NA, freq))) +
    geom_point() +
    scale_size(breaks = c(0.25, 0.50, 0.75, 1.00), limits = c(0, 1)) +
    scale_x_discrete(name = " ") +
    scale_y_discrete(name = "Signatures", limits = order$sig) +
    labs(
      color = "Mean relative \n attribution",
      size = "Proportion of tumors \n with the signature"
    ) +
    scale_color_viridis(limits = c(0, round(max(summary$perc_sig), 1))) +
    theme_minimal() +
    theme(
      legend.position = "right",
      panel.border = element_rect(colour = "black", fill = NA, size = .25),
      panel.grid.major = element_blank(),
      axis.text.x = element_blank()
    )
}

p <- cn_sig_landscape(data)
print(p)
ggsave(paste0("output/Figure_7a_", Sys.Date(), ".pdf"), device = "pdf", width = 3, height = 3.75)
